#include "pose_estimation.hpp"
#include <Eigen/Geometry>
#include <apriltag/apriltag_pose.h>
#include <apriltag/common/homography.h>
#include <opencv2/calib3d.hpp>
#include <tf2/convert.h>


geometry_msgs::msg::Transform
homography(apriltag_detection_t* const detection, const std::array<double, 4>& intr, double tagsize)
{
    apriltag_detection_info_t info = {detection, tagsize, intr[0], intr[1], intr[2], intr[3]};

    apriltag_pose_t pose;
    estimate_pose_for_tag_homography(&info, &pose);

    // rotate frame such that z points in the opposite direction towards the camera
    for(int i = 0; i < 3; i++) {
        // swap x and y axes
        std::swap(MATD_EL(pose.R, 0, i), MATD_EL(pose.R, 1, i));
        // invert z axis
        MATD_EL(pose.R, 2, i) *= -1;
    }

    return tf2::toMsg<apriltag_pose_t, geometry_msgs::msg::Transform>(const_cast<const apriltag_pose_t&>(pose));
}

geometry_msgs::msg::Transform
pnp(apriltag_detection_t* const detection, const std::array<double, 4>& intr, double tagsize)
{
    const std::vector<cv::Point3d> objectPoints{
        {-tagsize / 2, -tagsize / 2, 0},
        {+tagsize / 2, -tagsize / 2, 0},
        {+tagsize / 2, +tagsize / 2, 0},
        {-tagsize / 2, +tagsize / 2, 0},
    };

    const std::vector<cv::Point2d> imagePoints{
        {detection->p[0][0], detection->p[0][1]},
        {detection->p[1][0], detection->p[1][1]},
        {detection->p[2][0], detection->p[2][1]},
        {detection->p[3][0], detection->p[3][1]},
    };

    cv::Matx33d cameraMatrix;
    cameraMatrix(0, 0) = intr[0];// fx
    cameraMatrix(1, 1) = intr[1];// fy
    cameraMatrix(0, 2) = intr[2];// cx
    cameraMatrix(1, 2) = intr[3];// cy

    cv::Mat rvec, tvec;
    cv::solvePnP(objectPoints, imagePoints, cameraMatrix, {}, rvec, tvec);
    // // 调整 tvec 的轴数据
    // double x = tvec.at<double>(0, 0);// 原始 x
    // double y = tvec.at<double>(1, 0);// 原始 y
    // double z = tvec.at<double>(2, 0);// 原始 z

    // tvec.at<double>(0, 0) = z;// x 轴换成原始 y
    // tvec.at<double>(1, 0) = x;// y 轴换成原始 z 的负数
    // tvec.at<double>(2, 0) = y;// z 轴换成原始 x
    // // 将 rvec 转换为旋转矩阵
    // cv::Mat rotationMatrix;
    // cv::Rodrigues(rvec, rotationMatrix);

    // // 调整旋转矩阵的轴系
    // cv::Mat adjustedRotationMatrix = rotationMatrix.clone();
    // adjustedRotationMatrix.row(0) = rotationMatrix.row(1); // x 轴换成原始 y
    // adjustedRotationMatrix.row(1) = -rotationMatrix.row(2);// y 轴换成原始 z 的负数
    // adjustedRotationMatrix.row(2) = rotationMatrix.row(0); // z 轴换成原始 x

    // // 将调整后的旋转矩阵转换回旋转向量
    // cv::Rodrigues(adjustedRotationMatrix, rvec);

    return tf2::toMsg<std::pair<cv::Mat_<double>, cv::Mat_<double>>, geometry_msgs::msg::Transform>(std::make_pair(tvec, rvec));
}

const std::unordered_map<std::string, pose_estimation_f> pose_estimation_methods{
    {"homography", homography},
    {"pnp", pnp},
};
